"""
seq2seq模型
"""
import torch.nn as nn

from chatbot.encoder import Encoder
from chatbot.decoder import Decoder


class Seq2seq(nn.Module):
    def __init__(self):
        super(Seq2seq, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, input_data, input_length, target):
        encoder_output, _, encoder_hidden = self.encoder(input_data, input_length)
        decoder_output = self.decoder(target, encoder_hidden, encoder_output)

        return decoder_output

    def evaluate(self, input_data, input_length):
        encoder_output, _, encoder_hidden = self.encoder(input_data, input_length)
        pre = self.decoder.evaluate(encoder_hidden=encoder_hidden, encoder_output=encoder_output)

        return pre

    def evaluate_by_beamserch(self, input_data, input_length):
        encoder_output, _, encoder_hidden = self.encoder(input_data, input_length)
        pre = self.decoder.evaluate_by_beamsearch(encoder_hidden, encoder_output)

        return pre




